import tensorflow as tf
from tensorflow.python.keras.applications.densenet import DenseNet169, DenseNet121
from tensorflow.python.keras.applications.vgg16 import VGG16
from tensorflow.python.keras.callbacks import EarlyStopping
from tensorflow.python.keras.layers import GlobalAveragePooling2D, Dense, GlobalMaxPooling2D, MaxPooling2D, Flatten
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.optimizer_v2.adam import Adam

from util import read_data, get_data_new

AUTOTUNE = tf.data.experimental.AUTOTUNE

if __name__ == '__main__':
    train_generator, validation_generator = get_data_new()

    # 初始化DenseNet169网络(卷积神经网络的一种)
    mobile_net = VGG16(input_shape=(224, 224, 3), include_top=False)
    # 固定参数
    mobile_net.trainable = False

    model = Sequential([
        mobile_net,
        MaxPooling2D(),
        Flatten(),
        Dense(512, activation='relu'),
        Dense(4, activation='softmax')])
    adam = Adam(learning_rate=1e-4)
    model.compile(optimizer=adam,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])
    # print(model.summary())
    # steps_per_epoch = tf.math.ceil(len(all_image_paths) / BATCH_SIZE).numpy()

    early_stopping = EarlyStopping(
        monitor='accuracy',
        verbose=1,
        patience=10,
        mode='max',
        restore_best_weights=True
    )
    # # 迭代次数2000，准确率还可以，耐心等待
    history = model.fit(
        train_generator,  # 加载训练集的完整预处理部分
        callbacks=[early_stopping],
        steps_per_epoch=100,  # 每步训练加载100张图片
        epochs=2000,  # 模型训练10步
        validation_data=validation_generator,  # 加载测试集的完整预处理部分用于验证
        validation_steps=25,  # 每步加载25张测试集图片用于验证
        verbose=1)
